from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection
import random

# ===== 1. 连接到本地 Milvus =====
connections.connect(
    alias="default",
    host="127.0.0.1",  # 如果 Milvus 在 Docker 或远程运行，这里改成对应 IP
    port="19530"
)

print("✅ 已连接到 Milvus")

# ===== 2. 创建集合（如果不存在）=====
collection_name = "test_collection"

if utility.has_collection(collection_name):
    utility.drop_collection(collection_name)
    print(f"已删除旧集合: {collection_name}")

# 定义字段
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=128)  # 向量维度128
]

schema = CollectionSchema(fields, description="测试 Milvus 连接")

# 创建集合
collection = Collection(name=collection_name, schema=schema)
print(f"✅ 已创建集合: {collection_name}")

# ===== 3. 插入测试数据 =====
import numpy as np
num_vectors = 10
vectors = np.random.random((num_vectors, 128)).astype(np.float32)  # 10条128维随机向量

mr = collection.insert([vectors])
print(f"✅ 插入数据ID: {mr.primary_keys}")

# ===== 4. 创建索引（加快检索）=====
index_params = {
    "metric_type": "L2",  # L2距离
    "index_type": "IVF_FLAT",
    "params": {"nlist": 64}
}
collection.create_index(field_name="embedding", index_params=index_params)
print("✅ 已创建索引")

# ===== 5. 加载集合到内存 =====
collection.load()

# ===== 6. 随机选一个向量进行检索 =====
query_vector = [vectors[0]]
search_params = {"metric_type": "L2", "params": {"nprobe": 10}}

results = collection.search(
    data=query_vector,
    anns_field="embedding",
    param=search_params,
    limit=3,
    output_fields=["id"]
)

print("✅ 检索结果：")
for result in results[0]:
    print(f"  ID: {result.id}, 距离: {result.distance}")

# ===== 7. 断开连接 =====
connections.disconnect("default")
print("🔌 已断开连接")
